# -*- coding: utf-8 -*-
import pandas as pd
import numpy as np
import os
#显示所有列
pd.set_option('display.max_columns', None)
#显示所有行
pd.set_option('display.max_rows', None)
#设置value的显示长度为100，默认为50
pd.set_option('max_colwidth',100)

if __name__ == "__main__":
    base_dir = 'output'
    path = ['raw_sub_final_net_finetune_0.csv', 'raw_sub_final_1.csv', 'raw_sub_final_0.csv']
    res = []
    for p in path:
        df = pd.read_csv(os.path.join(base_dir, p))
        # print(df)
        df = df.astype({'pred_gender': str})
        for c in ['pred_age', 'pred_gender']:
            df[c] = df[c].apply(lambda x: np.array(list(map(float, x.split()))))
        res.append(df)
    sub = res[0]
    # print(sub.shape)
    for i in range(1, len(res)):
        sub = sub.merge(res[i], on='user_id')
        # print(sub.shape)
    # print(sub)
    cols = sub.columns
    sub['predicted_age'] = sub[[c for c in cols if 'pred_age' in c]].apply(lambda x: np.sum(x)/len(res), axis=1)
    sub['predicted_gender'] = sub[[c for c in cols if 'pred_gender' in c]].apply(lambda x: np.sum(x)/len(res), axis=1)
    sub = sub[['user_id', 'predicted_age', 'predicted_gender']]
    # print(sub)
    sub['predicted_age'] = sub['predicted_age'].apply(lambda x: np.argmax(x) + 1)
    # sub['predicted_gender'] = sub['predicted_gender'].apply(lambda x: 2 if x[0] > 0.5 else 1)
    sub['predicted_gender'] = sub['predicted_gender'].apply(lambda x: np.where(x > 0.5, 2, 1)[0])
    sub.to_csv('./output/final_merge.csv', index=False)




